package testrank.affinity;

import java.io.Serializable;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;

import testrank.TestRankConstants;

/**
 * @author Hagai Cibulski
 * 
 * contains set of non key words, and separate literals (strings and numbers)
 * 			this includes comments before the method
 * 
 * Example text in method:
 *	StringBuilder sayHello(String message) { 
 *		return new StringBuilder("hellooo!".subString(0, 5)); } // parameter is ignored
 *	} 
 */
public class WordGroup implements Serializable {

	private static final long serialVersionUID = 1L;

	// document represent methods
	// FUTURE represent classes?
	private static WordGroup documents = new WordGroup(); 

	private static int nDocuments = 0;

	private static double logDocuments = 0;
	
	/**
	 * English words from comments, simple identifiers, split identifiers (filter out keywords)
	 * These will be the input for wordDistance calculations 
	 * 
	 * Example: string, builder, say, hello, message, sub, parameter, ignored
	 */
	private SortedMap<String, Integer> words = new TreeMap<String, Integer>();
	
	/**
	 * Numbers, quoted strings, combined identifiers
	 * These will be the input for simple compare (maybe string edit distance; maybe more affinity propagation)
	 * 
	 * Example: StringBuilder, sayHello, hellooo!, subString, 0, 5
	 */
	private SortedMap<String, Integer> literals = new TreeMap<String, Integer>();
	
	private int wordsCount = 0;
	
	private int literalCount = 0;
	

	public boolean addWord(String term) {
		wordsCount++;
		return add(this.words, term);
	}

	public boolean addLiteral(String term) {
		literalCount++;
		return add(this.literals, term);
	}

	private boolean add(SortedMap<String, Integer> terms, String term) {
		Integer count = terms.get(term);
		if (count == null) {
			terms.put(term, 1);
			return true;
		}
		else {
			terms.put(term, count + 1);
			return false;
		}
	}

	public Set<String> getWords() {
		return this.words.keySet();
	}

	public Set<String> getLiterals() {
		return this.literals.keySet();
	}

	public double getWordFrequency(String term) {
		return getTermFrequency(words, term, wordsCount);
	}

	public double getLiteralFrequency(String term) {
		return getTermFrequency(literals, term, literalCount);
	}
	
	private double getTermFrequency(SortedMap<String, Integer> terms, String term, int termCount) {
		Integer count = terms.get(term);
		if (count == null || termCount == 0)
			return 0;
		
		return (double)count / (double)termCount;
	}

	public void addDocument() {
		nDocuments++;
		for (String term : getWords()) {
			documents.addWord(term);
		}
		for (String term : getLiterals()) {
			documents.addLiteral(term);
		}
	}
	
	public double getWordTFIDF(String term, Set<WordGroup> context) {
		return getTFIDF(term, words, wordsCount, documents.words, context, true);
	}
	
	public double getLiteralTFIDF(String term, Set<WordGroup> context) {
		return getTFIDF(term, literals, literalCount, documents.literals, context, false);
	}

	public double getTFIDF(String term, SortedMap<String, Integer> terms, int termsCount, SortedMap<String, Integer> documentsTerms, Set<WordGroup> context, boolean useWordAffinity) {
		if (nDocuments == 0)
			return 0;
		
		double tf = getTermFrequency(terms, term, termsCount);
		double idf = 
			TestRankConstants.AFFINITY_ALT_NORMALIZATION != TestRankConstants.AffinityNormalization.ByTfIdfContext ? 
				getInverseDocFrequency(term, documentsTerms) :
				getInverseDocFrequency(term, context, useWordAffinity);	
		double tfidf = tf * idf;
		assert tfidf >= 0;
		return tfidf;
	}


	private double getInverseDocFrequency(String term,
			SortedMap<String, Integer> documentsTerms) {
		double logD = logDocuments();
		Integer dn = documentsTerms.get(term);
		if (dn == null)
			return 0;
		
		double idf = logD - log(dn);
		return idf;
	}
	
	private double getInverseDocFrequency(String term,
			Set<WordGroup> context, boolean useWordAffinity) {
		double logD = log(context.size());
		
		int dn = 0;
		for (WordGroup g : context) {
			Set<String> docTerms = useWordAffinity ?
					g.getWords() : g.getLiterals();
					if (docTerms.contains(term)) {
						dn++;
			}
		}
		if (dn == 0)
			return 0;
		
		double idf = logD - log(dn);
		return idf;
	}


	private static double logDocuments() {
		if (logDocuments == 0)
			logDocuments = log(nDocuments);
		return logDocuments;
	}
	
	private static double log(double x) {
		return Math.log(x);
	}

	@Override
	public String toString() {
		return "WordGroup [literals=" + this.literals + ", words=" + this.words + "]";
	}
}
